import math, os, sys
from vtkpython import *



sample_atoms = (('C', (0,0,0)), ('C', (0,0,1)),
                ('H', (.5,0,1)), ('HO', (.5,.5,1)),
                ('O', (0,.5,.5)), ('Z',(1,1,1)))

sample_bonds = (('S', (0, 1)), ('S', (0, 3)), ('D', (1,4)))
                
                
carbon_bond = {'radius' : .05, 'color' : (.2,.2,.2) }
unknown_bond = {'radius' : .02, 'color' : (1, 1, 1) }
bondsByType = { 'S' : carbon_bond, 'unknown' : unknown_bond }
            
bondResolution = {'theta' : 20, 'phi' : 20 }


class BondShow:
  def __init__(self):
    self.polyData = None
    self.typeToIdx = None
    self.CreatePipeline()
    self.ConfigureBonds(bondsByType)
    
  def GetActor(self):
    return self.tubeActor
  
  def AddBonds(self, atomPts, bonds):
    """Set atom data
    
    atomSet - [(type, (x, y, z))] where type is C, H, O, HO.
    
    """
    
    self.polyData.SetPoints(atomPts)
    self.polyData.SetLines(bonds["cellArray"])
    self.polyData.GetCellData().SetScalars(bonds["radii"])
    self.polyData.GetCellData().AddArray(bonds["color"])

    self.RegisterPolydataInPipeline(self.polyData)

      
      
  def CreateAtomPoints(self, atomSet):
    createPts = vtkPoints()
    createPts.SetDataTypeToFloat()
    for atom in atomSet:
      ptIdx = createPts.InsertNextPoint(atom[1])
    return createPts

    
    
  def CreateCellArray(self, bondSet):
    createRadii = vtkFloatArray()
    createColor = vtkIdTypeArray()
    createColor.SetName("tubecolor")
    createCellArray = vtkCellArray()
    for (bondType, bondLink) in bondSet:
      if self.typeToIdx.has_key(bondType):
        typeIdx = self.typeToIdx[bondType]
        radius = bondsByType[bondType]["radius"]
      else:
        typeIdx = self.typeToIdx['unknown']
        radius = bondsByType['unknown']["radius"]
      createColor.InsertNextValue(typeIdx)
      createRadii.InsertNextValue(radius)
      createCellArray.InsertNextCell(2)
      createCellArray.InsertCellPoint(bondLink[0])
      createCellArray.InsertCellPoint(bondLink[1])
    
    return { "cellArray" : createCellArray, "color" : createColor,
             "radii" : createRadii}
      
      
      
  def AddLookupTableColor(self, index, color):
    self.lut.SetTableValue(index, color[0], color[1], color[2], 1)

  def CreatePipeline(self):
    self.polyData = vtk.vtkPolyData()
    self.tubeFilter = vtk.vtkTubeFilter()
    self.lut = vtk.vtkLookupTable()
    self.glyphMapper = vtk.vtkPolyDataMapper()
    self.tubeActor = vtk.vtkActor()

  def RegisterPolydataInPipeline(self, newPolyData):
    self.tubeFilter.SetInput(newPolyData)
    if self.polyData and self.polyData is not newPolyData:
      del self.polyData
    self.polyData = newPolyData

  def ConfigureBonds(self, bondTypes):
    self.typeToIdx = {}
    bondIdx = 0
    for bondType in bondTypes.keys():
      self.typeToIdx[bondType] = bondIdx
      bondIdx = bondIdx + 1
    bondCnt = bondIdx
    
    #self.tubeFilter.SetVaryRadiusToVaryRadiusByScalar()
    self.tubeFilter.SetVaryRadius(0)
    self.tubeFilter.SetRadius(0.05)
    self.tubeFilter.CappingOff()
    self.tubeFilter.SetNumberOfSides(12)
    self.tubeFilter.SetGenerateTCoordsToOff()
    
    self.lut.SetNumberOfTableValues(bondCnt)
    for sourceType in bondTypes.keys():
      idx = self.typeToIdx[sourceType]
      bondColor = bondTypes[sourceType]["color"]
      self.AddLookupTableColor(idx, bondColor)
    
    self.glyphMapper.SetInput(self.tubeFilter.GetOutput())
    self.glyphMapper.SetLookupTable(self.lut)
    self.glyphMapper.SetScalarRange(0,bondCnt-1)
    
    self.tubeActor.SetMapper(self.glyphMapper)
    self.tubeActor.GetProperty().SetOpacity(1.0)
    self.tubeActor.GetProperty().SetInterpolationToGouraud()
    
    
if __name__ == "__main__":
  actorCreator = BondShow()
  actorCreator.AddBonds(sample_atoms, sample_bonds)
  actorCreator.AddBonds(sample_atoms, sample_bonds)
  actorCreator.GetActor()